# -*- coding: utf-8; mode: snippet -*-
# name: Quantiles of Kaplan-Meier from lifelines
# key: kaplanmeier.quantiles@lifeline
# key: km_quantiles_lifeline
# contributor: Shuguang Sun
# --
def get_KM_times(survival_function,label):
    s = survival_function.reset_index()
    length = int(s.size / 2)
    s.rename(columns={label:"pred"}, inplace=True)
    below_75 = False
    below_50 = False
    below_25 = False
    estimates = [".",".","."]
    s["pred"] = s["pred"].apply(lambda x: round(x,6))
    for i in range(1,length):
        if (s["pred"][i] < 0.750000 and below_75 == False):
            if (s["pred"][i-1] == 0.750000):
                estimates[0] = (s["timeline"][i]+s["timeline"][i-1])/2
            else:
                estimates[0] = s["timeline"][i]
            below_75 = True
        if (s["pred"][i] < 0.500000 and below_50 == False):
            if (s["pred"][i-1] == 0.500000):
                estimates[1] = (s["timeline"][i]+s["timeline"][i-1])/2
            else:
                estimates[1] = s["timeline"][i]
            below_50 = True
        if (s["pred"][i] < 0.250000 and below_25 == False):
            if (s["pred"][i-1] == 0.250000):
                estimates[2] = (s["timeline"][i]+s["timeline"][i-1])/2
            else:
                estimates[2] = s["timeline"][i]
            below_25 = True
    return estimates